import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import scipy.stats as stats
import os
import warnings
import re
from pathlib import Path
from matplotlib.colors import ListedColormap
import matplotlib.patches as mpatches
warnings.filterwarnings('ignore')

# Set plotting style
plt.style.use('seaborn-v0_8')
sns.set_palette("colorblind")

def sanitize_filename(name):
    """Sanitize filenames by replacing invalid characters with underscores"""
    sanitized = re.sub(r'[<>:"/\\|?*]', '_', str(name))
    sanitized = re.sub(r'[ \t\n\r\f\v]', '_', sanitized)
    sanitized = re.sub(r'_+', '_', sanitized)
    return sanitized.strip('_')

def load_data_correctly(results_path):
    """Load data with proper parsing for comma-separated format"""
    print("Loading data...")
    df = pd.read_csv(results_path, sep=',')
    print(f"Successfully loaded data with shape: {df.shape}")
    return df

def extract_position(variant_id):
    """Extract position from variant ID, handling various formats safely"""
    try:
        parts = str(variant_id).split(':')
        if 'chr' in parts[0]:
            pos_part = parts[0].replace('chr', '').split('_')[0]
            return int(pos_part)
        else:
            return int(parts[0])
    except (ValueError, IndexError, AttributeError):
        return 1e9

def safe_log2_odds_ratio(odds_ratio):
    """Safely calculate log2 of odds ratio, handling zeros and negative values"""
    safe_or = odds_ratio.copy()
    safe_or[safe_or <= 0] = 1e-10
    return np.log2(safe_or)

def categorize_outcomes(all_results, p_threshold=0.05, min_variant_count=5):
    """
    Categorize outcomes into those with significant variants and those without
    """
    # Clean data first
    df_clean = all_results.copy()
    df_clean = df_clean.replace([np.inf, -np.inf], np.nan)
    df_clean = df_clean.dropna(subset=['P_Value', 'Odds_Ratio', 'Variant_Count'])
    
    # Apply filters
    significant_mask = (df_clean['P_Value'] < p_threshold) & (df_clean['Variant_Count'] >= min_variant_count)
    
    # Get outcomes with significant variants
    outcomes_with_significant = df_clean[significant_mask]['Outcome'].unique()
    all_outcomes = df_clean['Outcome'].unique()
    outcomes_without_significant = [outcome for outcome in all_outcomes if outcome not in outcomes_with_significant]
    
    print(f"Outcomes with significant variants: {len(outcomes_with_significant)}")
    print(f"Outcomes without significant variants: {len(outcomes_without_significant)}")
    
    return list(outcomes_with_significant), list(outcomes_without_significant)

def create_manhattan_subplot(ax, outcome_data, outcome_name, show_significance_lines=True):
    """
    Create a single Manhattan subplot for one outcome
    """
    if len(outcome_data) == 0:
        ax.text(0.5, 0.5, 'No data', transform=ax.transAxes, ha='center', va='center')
        ax.set_title(f"{outcome_name}\n(No data)", fontsize=10)
        return
    
    # Extract positions and calculate -log10 p-values
    outcome_data = outcome_data.copy()
    outcome_data['POS'] = outcome_data['Variant'].apply(extract_position)
    outcome_data = outcome_data[outcome_data['POS'] < 1e9]  # Filter invalid positions
    outcome_data['neg_log10_p'] = -np.log10(outcome_data['P_Value'])
    
    if len(outcome_data) == 0:
        ax.text(0.5, 0.5, 'No valid positions', transform=ax.transAxes, ha='center', va='center')
        ax.set_title(f"{outcome_name}\n(No valid positions)", fontsize=10)
        return
    
    # Sort by position
    outcome_data = outcome_data.sort_values('POS')
    
    # Color points by significance
    colors = ['red' if p < 0.05 else 'blue' for p in outcome_data['P_Value']]
    sizes = [60 if p < 0.05 else 30 for p in outcome_data['P_Value']]
    alphas = [0.8 if p < 0.05 else 0.5 for p in outcome_data['P_Value']]
    
    # Plot
    scatter = ax.scatter(outcome_data['POS'], outcome_data['neg_log10_p'],
                        c=colors, s=sizes, alpha=alphas, edgecolors='black', linewidth=0.3)
    
    # Add significance threshold lines
    if show_significance_lines:
        ax.axhline(y=-np.log10(0.05), color='red', linestyle='--', alpha=0.7, linewidth=1)
        ax.axhline(y=-np.log10(0.01), color='orange', linestyle='--', alpha=0.5, linewidth=0.8)
        ax.axhline(y=-np.log10(0.001), color='green', linestyle='--', alpha=0.5, linewidth=0.8)
    
    # Count significant variants for this outcome
    sig_count = len(outcome_data[outcome_data['P_Value'] < 0.05])
    total_count = len(outcome_data)
    
    # Set title with counts
    ax.set_title(f"{outcome_name}\n({sig_count}/{total_count} significant)", fontsize=10, pad=10)
    
    # Set labels
    ax.set_xlabel('Position', fontsize=8)
    ax.set_ylabel('-log10(p)', fontsize=8)
    
    # Format ticks
    ax.ticklabel_format(style='plain', axis='x')
    ax.tick_params(axis='both', which='major', labelsize=7)
    
    # Add grid
    ax.grid(True, alpha=0.3)
    
    return scatter

def create_comprehensive_manhattan_grid(all_results, outcomes_list, title_suffix, output_dir, 
                                       figsize=(20, 15), max_cols=4):
    """
    Create a grid of Manhattan plots for a list of outcomes
    """
    if len(outcomes_list) == 0:
        print(f"No outcomes to plot for {title_suffix}")
        return None
    
    # Calculate grid dimensions
    n_plots = len(outcomes_list)
    n_cols = min(max_cols, n_plots)
    n_rows = (n_plots + n_cols - 1) // n_cols  # Ceiling division
    
    # Create figure
    fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize)
    
    # Handle single subplot case
    if n_plots == 1:
        axes = np.array([axes])
    if n_rows == 1 and n_cols == 1:
        axes = np.array([axes])
    elif n_rows == 1:
        axes = axes.reshape(1, -1)
    elif n_cols == 1:
        axes = axes.reshape(-1, 1)
    
    # Flatten axes array for easy iteration
    axes_flat = axes.flatten()
    
    # Create subplots for each outcome
    for i, outcome in enumerate(outcomes_list):
        if i < len(axes_flat):
            ax = axes_flat[i]
            outcome_data = all_results[all_results['Outcome'] == outcome].copy()
            create_manhattan_subplot(ax, outcome_data, outcome)
        else:
            break
    
    # Hide unused subplots
    for i in range(len(outcomes_list), len(axes_flat)):
        axes_flat[i].set_visible(False)
    
    # Add main title
    plt.suptitle(f'Comprehensive Manhattan Plots - {title_suffix}', 
                fontsize=16, fontweight='bold', y=0.98)
    
    plt.tight_layout()
    
    # Save the plot
    filename = f"comprehensive_manhattan_grid_{sanitize_filename(title_suffix)}"
    plt.savefig(os.path.join(output_dir, f"{filename}.png"), 
                dpi=300, bbox_inches='tight', facecolor='white')
    plt.savefig(os.path.join(output_dir, f"{filename}.pdf"), 
                bbox_inches='tight', facecolor='white')
    plt.close()
    
    print(f"Created comprehensive Manhattan grid for {len(outcomes_list)} outcomes: {title_suffix}")
    return fig, axes

def create_summary_tables(all_results, outcomes_with_sig, outcomes_without_sig, output_dir):
    """
    Create summary tables for both significant and non-significant outcomes
    """
    # Summary for outcomes with significant variants
    sig_summary = []
    for outcome in outcomes_with_sig:
        outcome_data = all_results[all_results['Outcome'] == outcome].copy()
        sig_count = len(outcome_data[outcome_data['P_Value'] < 0.05])
        total_count = len(outcome_data)
        sig_summary.append({
            'Outcome': outcome,
            'Significant_Variants': sig_count,
            'Total_Variants': total_count,
            'Percentage_Significant': (sig_count / total_count * 100) if total_count > 0 else 0
        })
    
    sig_summary_df = pd.DataFrame(sig_summary)
    sig_summary_df = sig_summary_df.sort_values('Significant_Variants', ascending=False)
    sig_summary_df.to_csv(os.path.join(output_dir, "significant_outcomes_summary.csv"), index=False)
    
    # Summary for outcomes without significant variants
    non_sig_summary = []
    for outcome in outcomes_without_sig:
        outcome_data = all_results[all_results['Outcome'] == outcome].copy()
        total_count = len(outcome_data)
        non_sig_summary.append({
            'Outcome': outcome,
            'Total_Variants': total_count,
            'Significant_Variants': 0,
            'Percentage_Significant': 0
        })
    
    non_sig_summary_df = pd.DataFrame(non_sig_summary)
    non_sig_summary_df = non_sig_summary_df.sort_values('Total_Variants', ascending=False)
    non_sig_summary_df.to_csv(os.path.join(output_dir, "non_significant_outcomes_summary.csv"), index=False)
    
    print("Created summary tables for both outcome categories")
    return sig_summary_df, non_sig_summary_df

def create_legend_plot(output_dir):
    """
    Create a separate legend plot explaining the colors and symbols
    """
    fig, ax = plt.subplots(figsize=(8, 2))
    ax.axis('off')
    
    # Create legend elements
    legend_elements = [
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='red', markersize=10, label='Significant (p < 0.05)'),
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='blue', markersize=8, label='Non-significant (p ≥ 0.05)'),
        plt.Line2D([0], [0], color='red', linestyle='--', linewidth=2, label='p = 0.05 threshold'),
        plt.Line2D([0], [0], color='orange', linestyle='--', linewidth=1, label='p = 0.01 threshold'),
        plt.Line2D([0], [0], color='green', linestyle='--', linewidth=1, label='p = 0.001 threshold')
    ]
    
    ax.legend(handles=legend_elements, loc='center', frameon=True, fancybox=True, 
              shadow=True, ncol=3, fontsize=12)
    ax.set_title('Manhattan Plot Legend', fontsize=14, fontweight='bold', pad=20)
    
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "manhattan_plot_legend.png"), 
                dpi=300, bbox_inches='tight', facecolor='white')
    plt.savefig(os.path.join(output_dir, "manhattan_plot_legend.pdf"), 
                bbox_inches='tight', facecolor='white')
    plt.close()
    
    print("Created legend plot")

def main():
    # Define paths
    results_path = "/home/mahmo/TUMOR/latest/forest/last/forest_plots/all_association_results.csv"
    output_dir = "/home/mahmo/TUMOR/latest/forest/last/transformed_data/output"
    
    # Create output directory
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    print("Loading and processing results...")
    
    # Load data
    all_results = load_data_correctly(results_path)
    
    # Categorize outcomes into those with and without significant variants
    outcomes_with_significant, outcomes_without_significant = categorize_outcomes(all_results)
    
    print(f"\nOutcomes with significant variants ({len(outcomes_with_significant)}):")
    for outcome in outcomes_with_significant:
        outcome_data = all_results[all_results['Outcome'] == outcome]
        sig_count = len(outcome_data[outcome_data['P_Value'] < 0.05])
        total_count = len(outcome_data)
        print(f"  {outcome}: {sig_count} significant out of {total_count} total variants")
    
    print(f"\nOutcomes without significant variants ({len(outcomes_without_significant)}):")
    for outcome in outcomes_without_significant:
        outcome_data = all_results[all_results['Outcome'] == outcome]
        total_count = len(outcome_data)
        print(f"  {outcome}: {total_count} total variants (0 significant)")
    
    # Create comprehensive Manhattan grid for outcomes WITH significant variants
    if outcomes_with_significant:
        create_comprehensive_manhattan_grid(
            all_results, 
            outcomes_with_significant, 
            "Outcomes with Significant Variants", 
            output_dir,
            figsize=(20, 15),
            max_cols=4
        )
    
    # Create comprehensive Manhattan grid for outcomes WITHOUT significant variants
    if outcomes_without_significant:
        create_comprehensive_manhattan_grid(
            all_results, 
            outcomes_without_significant, 
            "Outcomes without Significant Variants", 
            output_dir,
            figsize=(20, 12),
            max_cols=4
        )
    
    # Create summary tables
    sig_summary_df, non_sig_summary_df = create_summary_tables(
        all_results, outcomes_with_significant, outcomes_without_significant, output_dir
    )
    
    # Create legend
    create_legend_plot(output_dir)
    
    # Save the categorized outcomes lists
    with open(os.path.join(output_dir, "outcome_categories.txt"), 'w') as f:
        f.write("OUTCOMES WITH SIGNIFICANT VARIANTS:\n")
        f.write("=" * 40 + "\n")
        for outcome in outcomes_with_significant:
            sig_count = len(all_results[(all_results['Outcome'] == outcome) & (all_results['P_Value'] < 0.05)])
            total_count = len(all_results[all_results['Outcome'] == outcome])
            f.write(f"{outcome}: {sig_count} significant out of {total_count} total variants\n")
        
        f.write("\nOUTCOMES WITHOUT SIGNIFICANT VARIANTS:\n")
        f.write("=" * 40 + "\n")
        for outcome in outcomes_without_significant:
            total_count = len(all_results[all_results['Outcome'] == outcome])
            f.write(f"{outcome}: {total_count} total variants (0 significant)\n")
    
    print(f"\nAnalysis complete! All plots and data saved to: {output_dir}")
    print(f"Created {len(outcomes_with_significant)} Manhattan plots for outcomes WITH significant variants")
    print(f"Created {len(outcomes_without_significant)} Manhattan plots for outcomes WITHOUT significant variants")

if __name__ == "__main__":
    main()